"""
PyTorch数据集类
用于加载图片和对应的标签数据
"""

import os
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import torchvision.transforms as transforms
from label_mapping import LABEL_TO_ID, get_label_description
import pandas as pd

class ESFDataset(Dataset):
    """ES和扩展F信号检测数据集"""
    
    def __init__(self, image_dir, labels_file, image_names_file, transform=None, train=True):
        """
        Args:
            image_dir (str): 图片文件夹路径
            labels_file (str): 标签文件路径 (0818.txt)
            image_names_file (str): 图片名称文件路径 (mingzi.txt)
            transform: 图像变换
            train (bool): 是否为训练模式
        """
        self.image_dir = image_dir
        self.transform = transform
        self.train = train
        
        # 读取标签和图片名称
        with open(labels_file, 'r', encoding='utf-8') as f:
            self.labels = [line.strip() for line in f.readlines()]
            
        with open(image_names_file, 'r', encoding='utf-8') as f:
            self.image_names = [line.strip() for line in f.readlines()]
            
        # 确保标签和图片名称数量一致
        assert len(self.labels) == len(self.image_names), \
            f"标签数量({len(self.labels)})与图片数量({len(self.image_names)})不匹配"
        
        # 过滤掉不存在的图片
        self.valid_indices = []
        for i, img_name in enumerate(self.image_names):
            img_path = os.path.join(image_dir, img_name)
            if os.path.exists(img_path):
                self.valid_indices.append(i)
        
        print(f"总共 {len(self.image_names)} 个样本，有效样本 {len(self.valid_indices)} 个")
        
        # 统计标签分布
        self.print_label_distribution()
    
    def __len__(self):
        return len(self.valid_indices)
    
    def __getitem__(self, idx):
        # 获取有效索引
        real_idx = self.valid_indices[idx]
        
        # 获取图片路径和标签
        img_name = self.image_names[real_idx]
        img_path = os.path.join(self.image_dir, img_name)
        label_str = self.labels[real_idx]
        
        # 加载图片
        try:
            image = Image.open(img_path).convert('RGB')
        except Exception as e:
            print(f"无法加载图片 {img_path}: {e}")
            # 返回一个默认的黑色图片
            image = Image.new('RGB', (224, 224), (0, 0, 0))
        
        # 应用变换
        if self.transform:
            image = self.transform(image)
        
        # 转换标签为数字
        label = LABEL_TO_ID.get(label_str, -1)
        if label == -1:
            print(f"警告: 未知标签 '{label_str}' 在索引 {real_idx}")
            label = 0  # 默认为第一个类别
        
        return image, label, img_name
    
    def print_label_distribution(self):
        """打印标签分布"""
        from collections import Counter
        valid_labels = [self.labels[i] for i in self.valid_indices]
        label_counts = Counter(valid_labels)
        
        print("\n标签分布:")
        print("-" * 60)
        for label, count in label_counts.most_common():
            desc = get_label_description(label)
            percentage = count / len(valid_labels) * 100
            print(f"{label:6} ({desc:15}) : {count:4} 个 ({percentage:5.1f}%)")
        print("-" * 60)

def get_transforms(image_size=224, train=True):
    """获取图像变换"""
    if train:
        return transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomRotation(degrees=10),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
    else:
        return transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])

def create_data_loaders(image_dir, labels_file, image_names_file, 
                       batch_size=32, train_split=0.8, image_size=224):
    """创建训练和验证数据加载器"""
    
    # 创建完整数据集
    full_dataset = ESFDataset(
        image_dir=image_dir,
        labels_file=labels_file,
        image_names_file=image_names_file,
        transform=get_transforms(image_size, train=True)
    )
    
    # 分割训练和验证集
    total_size = len(full_dataset)
    train_size = int(train_split * total_size)
    val_size = total_size - train_size
    
    train_dataset, val_dataset = torch.utils.data.random_split(
        full_dataset, [train_size, val_size],
        generator=torch.Generator().manual_seed(42)
    )
    
    # 为验证集设置不同的变换
    val_transform = get_transforms(image_size, train=False)
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True, 
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=4,
        pin_memory=True
    )
    
    print(f"\n数据集分割:")
    print(f"训练集: {train_size} 个样本")
    print(f"验证集: {val_size} 个样本")
    
    return train_loader, val_loader

if __name__ == "__main__":
    # 测试数据集
    dataset = ESFDataset(
        image_dir=".",
        labels_file="0818.txt",
        image_names_file="mingzi.txt",
        transform=get_transforms()
    )
    
    print(f"\n数据集大小: {len(dataset)}")
    
    # 测试加载一个样本
    if len(dataset) > 0:
        image, label, img_name = dataset[0]
        print(f"图片形状: {image.shape}")
        print(f"标签: {label}")
        print(f"图片名称: {img_name}")